# You need to set the environmental variable OPENAI_API_KEY
from langchain_core.tools import tool
from langchain import hub
from langchain_zhipu import ChatZhipuAI

from langchain.agents import create_react_agent
from langchain.agents import AgentType
from langchain.agents import AgentExecutor

from langchain_core.prompts import PromptTemplate

from logger import setup_logger
from langchain.agents.output_parsers import ReActJsonSingleInputOutputParser
from langchain.tools.render import render_text_description_and_args

logger = setup_logger()


# Custom tool for the Agent
@tool
def get_employee_id(name):
    """
    Obtain employee ID based on employee name, receive value such as Evan, return E008
    """

    logger.info("getEmployeeId:" + name)
    fake_employees = {
        "Alice": "E001",
        "Bob": "E002",
        "Charlie": "E003",
        "Diana": "E004",
        "Evan": "E005",
        "Fiona": "E006",
        "George": "E007",
        "Hannah": "E008",
        "Ian": "E009",
        "Jasmine": "E010"}

    return fake_employees.get(name, "Employee not found")


# Custom tool for the Agent
@tool
def get_employee_salary(employee_id):
    """
    To get the salary of an employee, it takes employee_id as input and return salary receive value such as E008, return 53000
    """

    logger.info("get_employee_salary:" + employee_id)
    employee_salaries = {
        "E001": 56000,
        "E002": 47000,
        "E003": 52000,
        "E004": 61000,
        "E005": 45000,
        "E006": 58000,
        "E007": 49000,
        "E008": 53000,
        "E009": 50000,
        "E010": 55000
    }
    return employee_salaries.get(employee_id, "Employee not found")


# Saved React Prompt in langchain hub, we could manually type the prompt as well.
# prompt = hub.pull("hwchase17/react")

template = '''
SYSTEM

Answer the following questions as best you can. You have access to the following tools:



{tools}



The way you use the tools is by specifying a json blob.

Specifically, this json should have a `action` key (with the name of the tool to use) and a `action_input` key (with the input to the tool going here).



The only values that should be in the "action" field are: {tool_names}



The $JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. Here is an example of a valid $JSON_BLOB:



```

{{

  "action": $TOOL_NAME,

  "action_input": $INPUT

}}

```



ALWAYS use the following format:



Question: the input question you must answer

Thought: you should always think about what to do

Action:

```

$JSON_BLOB

```

Observation: the result of the action

... (this Thought/Action/Observation can repeat N times)

Thought: I now know the final answer

Final Answer: the final answer to the original input question



Begin! Reminder to always use the exact characters `Final Answer` when responding.

HUMAN

{input}



{agent_scratchpad}'''

prompt = PromptTemplate.from_template(template)
model = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                    , model="glm-4", verbose=True)
model.do_sample = False  # 温度设置为0，结果随机性 ghbnm

tools = [get_employee_salary, get_employee_id]
agent = create_react_agent(
    model,
    tools,
    prompt,
    tools_renderer=render_text_description_and_args,
    output_parser=ReActJsonSingleInputOutputParser(),
)
# ,max_iterations=2
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, handle_parsing_errors=True)
answer = agent_executor.invoke({"input": "Alice的工资是多少？?"})
print(answer)
